import cv2
import csv
import os
import random

# 定义每段的帧数
frames_clip = [17, 51, 102]
output_dir = "/g0022010zlc/zyz/varia_data"
output_csv = "/g0022010zlc/zyz/varia.csv"
# has_done = []
to_do = {}
# with open(output_csv, newline='') as csvfile:
#     reader = csv.DictReader(csvfile)
#     for row in reader:
#         video_path = row["path"] 
#         if not os.path.exists(video_path):
#             rel = "_".join(row['relpath'].split("_")[:-1]) + ".mp4"
#             if rel not in to_do.keys():
#                 to_do[rel] = int(row["num_frames"])
# print("to do len:", len(to_do))
# with open(output_csv, newline='') as csvfile:
#     reader = csv.DictReader(csvfile)
#     for row in reader:
#         video_path = "_".join(row['relpath'].split("_")[:-1]) + ".mp4"
#         if video_path not in has_done:
#             has_done.append(video_path)
# print(has_done)
# 读取 CSV 文件
csv_file = '/g0022010zlc/zyz/opensora_distill/origin.csv'
with open(csv_file, newline='') as csvfile:
    reader = csv.DictReader(csvfile)
    for row in reader:
        video_path = row['path']
        # relpath = row["relpath"]
        # if relpath not in to_do.keys():
        #     print(video_path + " done!:")
        #     continue
        fps = row["fps"]
        total_frames = row['num_frames']
        video_path_split = video_path.split("/")
        sub1, sub2 = video_path_split[-3], video_path_split[-2]
        video_name, video_ext = os.path.splitext(video_path_split[-1])
        
        # 打开视频文件
        cap = cv2.VideoCapture(video_path)
        
        # 检查视频是否成功打开
        if not cap.isOpened():
            print(f"Error: Could not open video {video_path}.")
            continue
        
        clip_count = 0
        # frames_per_clip = to_do[relpath]
       
        # 读取视频帧并保存每段视频

        for i in range(4):
            p = random.random()
            if p < 0.4:
                frames_per_clip = frames_clip[0]
            elif p < 0.9:
                frames_per_clip = frames_clip[1]
            else:
                frames_per_clip = frames_clip[2]
            frames = []
            if int(total_frames) < frames_per_clip:
                continue
            # 随机选择一个起始帧
            start_frame = random.randint(0, int(total_frames) - frames_per_clip - 1)
            cap.set(cv2.CAP_PROP_POS_FRAMES, start_frame)
            for _ in range(frames_per_clip):
                ret, frame = cap.read()
                if not ret:
                    break
                frames.append(frame)
            
            # 保存当前视频片段
            clip_count += 1
            if not os.path.exists(f"{output_dir}/{sub1}/{sub2}"):
                os.makedirs(f"{output_dir}/{sub1}/{sub2}")
            output_path = f"{output_dir}/{sub1}/{sub2}/{video_name}_{clip_count}{video_ext}"
            fourcc = cv2.VideoWriter_fourcc(*'mp4v')
            out = cv2.VideoWriter(output_path, fourcc, int(float(fps)), (frames[0].shape[1], frames[0].shape[0]))
            
            for frame in frames:
                out.write(frame)
            out.release()
            print(f"Saved {video_name}")
            row["path"] = output_path
            row["id"] = f"{video_name}_{clip_count}"
            row["relpath"] = f"{sub1}/{sub2}/{video_name}_{clip_count}{video_ext}"
            row["num_frames"] =frames_per_clip 
            with open(output_csv, 'a', newline='') as file:
                fieldnames = ["path","id","relpath","num_frames","height","width","aspect_ratio","fps","resolution","text"]
                writer = csv.DictWriter(file, fieldnames=fieldnames)
                writer.writerow(row)
        
        # 释放视频捕获对象
        cap.release()
        print(f"Finished processing {video_path}")
print("All videos processed.")
